#include <fstream>
#include "gc.h"
#include "itkImageRegionConstIterator.h"
#include "itkImageRegionConstIteratorWithIndex.h"

typedef std::map<LabelPixel, int> LabelCountMap;

void find_initial (LabelPixel truthLabel, std::string truthImageName, 
		   std::string waterImageName)
{
  LabelImage::Pointer truthImage = readLabelImage(truthImageName.c_str());
  LabelImage::Pointer waterImage = readLabelImage(waterImageName.c_str());
  LabelCountMap labelCountMap;
  for (itk::ImageRegionConstIteratorWithIndex<LabelImage> 
	 itr(truthImage, truthImage->GetLargestPossibleRegion()); 
       !itr.IsAtEnd(); ++itr) {
    if (itr.Get() == truthLabel) {
      LabelPixel waterLabel = waterImage->GetPixel(itr.GetIndex());
      if (labelCountMap.find(waterLabel) != labelCountMap.end()) {
	labelCountMap[waterLabel]++;
      }
      else {
	labelCountMap[waterLabel] = 1;
      }
    }
  }
  for (LabelCountMap::const_iterator itr = labelCountMap.begin(); 
       itr != labelCountMap.end(); itr++) {
    int waterArea = 0;
    for (itk::ImageRegionConstIteratorWithIndex<LabelImage> 
	   iitr(waterImage, waterImage->GetLargestPossibleRegion()); 
	 !iitr.IsAtEnd(); ++iitr) {
      if (iitr.Get() == itr->first) {
	waterArea++;
      }
    }
    std::cout << "Label " << itr->first << ": " 
	      << double(itr->second) / double(waterArea) 
	      << "(" << itr->second << "/" << waterArea << ")" << std::endl;
  }
}



void find_initial2 (LabelPixel truthLabel, std::string truthImageName, 
		    std::string waterImageName)
{
  LabelImage::Pointer truthImage = readLabelImage(truthImageName.c_str());
  LabelImage::Pointer waterImage = readLabelImage(waterImageName.c_str());
  LabelPixel maxWaterLabel = 0;
  for (itk::ImageRegionConstIterator<LabelImage> 
	 itr(waterImage, waterImage->GetLargestPossibleRegion());
       !itr.IsAtEnd(); ++itr) {
    if (itr.Get() > maxWaterLabel) {
      maxWaterLabel = itr.Get();
    }
  }
  for (LabelPixel waterLabel = 1; waterLabel <= maxWaterLabel; waterLabel++) {
    Component component;
    for (itk::ImageRegionConstIteratorWithIndex<LabelImage> 
	   itr(waterImage, waterImage->GetLargestPossibleRegion()); 
	 !itr.IsAtEnd(); ++itr) {
      if (itr.Get() == waterLabel) {
	Point point(itr.GetIndex());
	component.points.push_back(point);
      }
    }
    LabelImage::IndexType center;
    int xc, yc;
    computeComponentInsideCenter(waterImage, component, xc, yc, VERBOSE_NONE);
    center[0] = xc;
    center[1] = yc;
    if (truthImage->GetPixel(center) == truthLabel) {
      std::cout << "Water label: " << waterLabel << std::endl;
    }
  }
}



// Find corresponding watershed components for each truth component
void find_initial3 (std::string truthImageName, std::string waterImageName, 
		    std::string outputInitialFileName)
{
  LabelImage::Pointer truthImage = readLabelImage(truthImageName.c_str());
  LabelImage::Pointer waterImage = readLabelImage(waterImageName.c_str());
  LabelPixel maxTruthComponentNo = 0;
  LabelPixel maxWaterComponentNo = 0;
  for (itk::ImageRegionConstIterator<LabelImage> 
	 itr(truthImage, truthImage->GetLargestPossibleRegion()); 
       !itr.IsAtEnd(); ++itr) {
    if (itr.Get() > maxTruthComponentNo) {
      maxTruthComponentNo = itr.Get();
    }
  }
  for (itk::ImageRegionConstIterator<LabelImage> 
	 itr(waterImage, waterImage->GetLargestPossibleRegion()); 
       !itr.IsAtEnd(); ++itr) {
    if (itr.Get() > maxWaterComponentNo) {
      maxWaterComponentNo = itr.Get();
    }
  }
  int truthComponentNum = int(maxTruthComponentNo) + 1;
  int waterComponentNum = int(maxWaterComponentNo) + 1;
  std::vector<std::vector<int> > componentNoSets(truthComponentNum);
  for (int waterComponentNo = 0; waterComponentNo < waterComponentNum; 
       waterComponentNo++) {
    LabelCountMap labelCountMap; // area of truth labels
    std::vector<int> labels; // truth labels
    for (itk::ImageRegionConstIteratorWithIndex<LabelImage> 
	   itr(waterImage, waterImage->GetLargestPossibleRegion()); 
	 !itr.IsAtEnd(); ++itr) {
      if (itr.Get() == LabelPixel(waterComponentNo)) {
        int  truthComponentNo = truthImage->GetPixel(itr.GetIndex());
	if (labelCountMap.count(truthComponentNo) > 0) {
	  labelCountMap[truthComponentNo]++;
	}
	else {
	  labelCountMap[truthComponentNo] = 1;
	  labels.push_back(truthComponentNo);
	}
      }
    }
    int maxCountLabel = 0;	// max area truth label
    int maxCount = 0;		// max area
    for (std::vector<int>::const_iterator itr = labels.begin(); 
	 itr != labels.end(); itr++) {
      if (labelCountMap[*itr] > maxCount) {
	maxCount = labelCountMap[*itr];
	maxCountLabel = *itr;
      }
    }
    componentNoSets[maxCountLabel].push_back(waterComponentNo);
  }
  std::ofstream fs(outputInitialFileName.c_str());
  for (int truthComponentNo = 0; truthComponentNo < truthComponentNum; 
       truthComponentNo++) {
    fs << truthComponentNo << " " << componentNoSets[truthComponentNo].size();
    for (std::vector<int>::const_iterator itr = 
	   componentNoSets[truthComponentNo].begin(); 
	 itr != componentNoSets[truthComponentNo].end(); itr++) {
      fs << " " << *itr;
    }
    fs << std::endl;
  }
  fs.close();
}



int main (int argc, char* argv[])
{
  int argNo = 1;
  // LabelPixel truthLabel = atoi(argv[argNo++]);
  std::string outputInitialFileName = argv[argNo++];
  std::string truthImageName = 
    // "/home/ting/Workspace/data/truth/pos_flood/label002_slice200_pos_flood_new.mha";
    "/home/sci/ting/Workspace/data/truth/pos_flood/label002_slice200_pos_flood_new.mha";
  std::string waterImageName = 
    // "/home/ting/Workspace/data/water/07/slice_200_cv8_tv_blur_water_07.png";
    "/home/sci/ting/Workspace/data/water/07/slice_200_cv8_tv_blur_water_07.png";
  // find_initial(truthLabel, truthImageName, waterImageName);
  // find_initial2(truthLabel, truthImageName, waterImageName);
  find_initial3(truthImageName, waterImageName, outputInitialFileName);
  return 0;
}
